import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# 定义马尔可夫链的转移矩阵和状态
transition_matrix = np.array([
    [0.7, 0.3],  # A -> A, A -> B
    [0.4, 0.6]   # B -> A, B -> B
])
states = ['A', 'B']
state_indices = {state: idx for idx, state in enumerate(states)}

# 初始化当前状态
current_state = np.random.choice(states)

# 构建有向图用于可视化
G = nx.DiGraph()

# 添加节点
for state in states:
    G.add_node(state)

# 添加边及其权重
edges = [
    ('A', 'A', {'weight': transition_matrix[0][0]}),
    ('A', 'B', {'weight': transition_matrix[0][1]}),
    ('B', 'A', {'weight': transition_matrix[1][0]}),
    ('B', 'B', {'weight': transition_matrix[1][1]})
]
G.add_edges_from(edges)

# 布局设置
pos = {
    'A': (0, 0),
    'B': (1, 0)
}

# 动画参数
num_steps = 50
path = [current_state]

# 更新函数，用于动画每一帧
def update(frame):
    global current_state
    current_idx = state_indices[current_state]
    next_state = np.random.choice(states, p=transition_matrix[current_idx])
    path.append(next_state)
    current_state = next_state

    # 清空并重新绘制
    ax.clear()
    ax.set_title(f"Step {frame + 1}: Current State = {current_state}")

    # 绘制状态转移图
    nx.draw(G, pos, with_labels=True, node_size=800, node_color='lightblue', font_size=16, ax=ax, arrows=True)

    # 标注当前状态
    nx.draw_networkx_nodes(G, pos, nodelist=[current_state], node_color='orange', node_size=900, ax=ax)

    # 显示路径
    print(f"Step {frame + 1}: Path = {' -> '.join(path)}")

    return []

# 创建绘图对象
fig, ax = plt.subplots(figsize=(6, 4))
ani = FuncAnimation(fig, update, frames=num_steps, interval=1000, blit=True)

plt.show()

# 输出最终路径
print("\nFinal Markov Chain Path:")
print(" -> ".join(path))
